#!/usr/bin/env python3
#
# @file MainApp.py
# @brief
# @author QRS
# @version 1.0
# @date 2024-04-03 11:34

import cv2
import numpy as np
import tkinter as tk
from PIL import Image, ImageTk

from Utils import Logger
from HiUsbCamera import HiCamera
from OnnxInfer import model_infer, input_size

UPDATE_INTERVAL_MS = 3000

ROI_POINTS = {
    'DA2681354': [
        (662, 343),
        (2784, 455),
        (714, 737),
        (2814, 842),
        (708, 1289),
        (2854, 1368),
        (560, 2045),
        (3170, 2144)
    ],
}
    # (1118, 258),
    # (2524, 304),
    # (1078, 560),
    # (2656, 639),
    # (842, 987),
    # (2978, 1092),
    # (454, 1582),
    # (3477, 1690)

INPUT_W, INPUT_H = input_size
L_PAD_SIZE, R_PAD_SIZE = 20, 20
ROI_PAD_SIZE = L_PAD_SIZE + R_PAD_SIZE
PAD_VALUE = 0
NROW = 4

l_padding = np.ones((INPUT_H, L_PAD_SIZE, 3), dtype=np.uint8) * PAD_VALUE
r_padding = np.ones((INPUT_H, R_PAD_SIZE, 3), dtype=np.uint8) * PAD_VALUE
t_padding = np.ones((L_PAD_SIZE, INPUT_W + ROI_PAD_SIZE, 3), dtype=np.uint8) * PAD_VALUE
b_padding = np.ones((R_PAD_SIZE, INPUT_W + ROI_PAD_SIZE, 3), dtype=np.uint8) * PAD_VALUE


class SimpleApp(object):
    def __init__(self, window, window_title, win_width, win_height, cameras):
        self.window = window
        self.window.title(window_title)
        self.cameras = cameras

        Logger.info(f'Canvas width: {win_width}, height: {win_height}')
        self.canvas = tk.Canvas(window, width=win_width, height=win_height)
        self.canvas.pack()
        self.update_image()
        self.window.mainloop()

    def update_image(self):
        images = []
        for cam in self.cameras:
            frame = cam.read()
            images.extend([frame[y1:y2, x1:x2] for x1, y1, x2, y2 in cam.roi_boxes])
        preds = model_infer(images, debug=True)
        for i in range(preds.shape[0]):
            pred_mask_contours = cv2.findContours(preds[i], cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[0]
            if len(pred_mask_contours) > 0:
                max_cont = max(pred_mask_contours, key=cv2.contourArea)
                cv2.drawContours(images[i], [max_cont], 0, (0, 0, 255), 2)
        h_imgs = []
        for i, img in enumerate(images[:16]):
            h_imgs.append(np.vstack((t_padding, np.hstack((l_padding, img, r_padding)), b_padding)))
        v_imgs = []
        for i in range(0, len(h_imgs), NROW):
            v_imgs.append(np.hstack(h_imgs[i:i + NROW]))
        grid_image = np.vstack(v_imgs)
        Logger.info(f'Grid image size: {grid_image.shape}')
        self.photo = ImageTk.PhotoImage(Image.fromarray(grid_image))
        self.canvas.create_image(0, 0, image=self.photo, anchor=tk.NW)
        self.window.after(UPDATE_INTERVAL_MS, self.update_image)

def app_main():
    offset = input_size[0] // 2
    cameras = []
    for dinfo in HiCamera.list():
        did = HiCamera.get_serial_number(dinfo)
        if did in ROI_POINTS:
            boxes8 = []
            for x, y in ROI_POINTS[did]:
                 boxes8.append((x - offset, y - offset, x + offset, y + offset))
            cam = HiCamera(did).open(dinfo)
            cam.roi_boxes = boxes8
            cameras.append(cam)

    if len(cameras) == 0:
        Logger.warn('Not found valid usb cameras')
        return

    win_width, win_height = (input_size[0] + ROI_PAD_SIZE) * NROW, (input_size[1] + ROI_PAD_SIZE) * NROW

    SimpleApp(tk.Tk(), 'WeldingSpotsSeg', win_width, win_height, cameras)

if __name__ == '__main__':
    app_main()
